import triton
import triton.language as tl
import test_common
import pytest
import torch
import torch_npu
import numpy as np
import random

@triton.jit
def kernel_softmax_online(
    x_ptr,
    x_row_stride,
    y_ptr,
    y_row_stride,
    n_rows,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    row_start = tl.program_id(0)
    row_step = tl.num_programs(0)
    for row_idx in tl.range(row_start, n_rows, row_step):
        x_ptr_base = x_ptr + row_idx * x_row_stride
        y_ptr_base = y_ptr + row_idx * y_row_stride

        mm = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - float("inf")
        ss = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
        for i in range(0, tl.cdiv(n_cols, BLOCK_SIZE)):
            idx = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
            x = tl.load(x_ptr_base + idx, mask=idx < n_cols, other=-float("inf"))
            mm_new = tl.maximum(mm, x)
            if i:  # 第 1 轮不需要，且容易整出 nan
                ss *= tl.exp(mm - mm_new)
            x = tl.exp(x - mm_new)
            ss += tl.where(idx < n_cols, x, 0.0)
            mm = mm_new

        mm_new = tl.max(mm)
        ss *= tl.exp(mm - mm_new)
        ss = tl.sum(ss)
        mm = mm_new

        eps = float(1e-9)
        ss = tl.maximum(ss, eps)

        for i in range(0, tl.cdiv(n_cols, BLOCK_SIZE)):
            idx = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
            x = tl.load(x_ptr_base + idx, mask=idx < n_cols, other=-float("inf"))
            x = tl.exp(x - mm) / ss
            tl.store(y_ptr_base + idx, x, mask=idx < n_cols)


def naive_softmax(x):
    # read  MN elements ; write M  elements
    x_max = x.max(dim=1)[0]
    # read MN + M elements ; write MN elements
    z = x - x_max[:, None]
    # read  MN elements ; write MN elements
    numerator = torch.exp(z)
    # read  MN elements ; write M  elements
    denominator = numerator.sum(dim=1)
    # read MN + M elements ; write MN elements
    ret = numerator / denominator[:, None]
    # in total: read 5MN + 2M elements ; wrote 3MN + 2M elements
    return ret


def softmax(x):
    n_rows, n_cols = x.shape
    # BLOCK_SIZE = triton.next_power_of_2(n_cols)
    y = torch.empty_like(x).npu()
    num_programs = 40
    num_programs = min(num_programs, n_rows)
    kernel_softmax_online[(num_programs, 1, 1)](
        x,
        x.stride(0),
        y,
        y.stride(0),
        n_rows,
        n_cols,
        BLOCK_SIZE=32
    )
    return y

param_lists = [
    [torch.float32, "float32", (1823, 781)],
    [torch.float16, "float16", (1823, 100)],
    [torch.bfloat16, "bfloat16", (1823, 256)],
]

@pytest.mark.parametrize('param_lists',param_lists)
def test_softmax(param_lists):
    dtype, sigtype, shape = param_lists
    torch.manual_seed(42)
    np.random.seed(42) 
    random.seed(42)
    torch.npu.set_device(7)

    x = test_common.generate_tensor(shape, sigtype).npu()
    y_triton = softmax(x)
    y_torch = naive_softmax(x)
    test_common.validate_cmp(sigtype, y_triton.cpu(), y_torch.cpu())